{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "da3a8407-0237-4ece-b268-75a396258764",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import os.path as osp\n",
    "import fiftyone as fo\n",
    "import fiftyone.zoo as foz\n",
    "import fiftyone.brain as fob\n",
    "\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "125b02e5-9fde-41ac-9d71-ab266b5f80da",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_name=\"nike\"\n",
    "fo.delete_datasets('*')\n",
    "\n",
    "dataset = fo.Dataset(name=dataset_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "1f193fa6-6c75-4ba3-89d4-46eaed494682",
   "metadata": {},
   "outputs": [],
   "source": [
    "root_dir = '/data/pt/data/all/shoe_dataset'\n",
    "\n",
    "split = \"val\"\n",
    "img_dir = osp.join(root_dir, \"images\",split)\n",
    "anno_dir = osp.join(root_dir, \"labels\",split)\n",
    "\n",
    "preds_dir = \"/home/user/workspace/deep_learning_tools/ultralytics/runs/detect/val3/labels\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "f91444fd-91ad-42bf-88ef-5fd6b097fcc3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def bbox_from_polygon(points):\n",
    "    \"\"\"\n",
    "    输入：points 是 Nx2 的 numpy 数组，表示多边形顶点 [[x1, y1], [x2, y2], ...]\n",
    "    输出：x_min, y_min, x_max, y_max\n",
    "    \"\"\"\n",
    "    points = np.array(points)\n",
    "    x_min = np.min(points[:, 0])\n",
    "    y_min = np.min(points[:, 1])\n",
    "    x_max = np.max(points[:, 0])\n",
    "    y_max = np.max(points[:, 1])\n",
    "    return x_min, y_min, x_max, y_max"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "c04ed43e-ec93-4e77-8479-42a7667ca76b",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Anomaly:\n",
    "    def __init__(self, score, bbox, polygon, category):\n",
    "        \"\"\"\n",
    "        :param score: float，缺陷置信度\n",
    "        :param bbox: list of 4 float，[x_min, y_min, x_max, y_max] 或其他格式\n",
    "        :param polygon: list of (x, y) tuples，多边形坐标点\n",
    "        :param category: str 或 int，类别名称或类别 ID\n",
    "        \"\"\"\n",
    "        self.score = score\n",
    "        self.bbox = bbox\n",
    "        self.polygon = polygon\n",
    "        self.category = category\n",
    "\n",
    "    def __repr__(self):\n",
    "        return (\n",
    "            f\"Anomaly(score={self.score:.2f}, \"\n",
    "            f\"bbox={self.bbox}, \"\n",
    "            f\"polygon_len={len(self.polygon)}, \"\n",
    "            f\"category={self.category})\"\n",
    "        )\n",
    "class BBox:\n",
    "    def __init__(self, x1, y1, x2, y2):\n",
    "        \"\"\"\n",
    "        使用 xyxy 格式初始化边界框。\n",
    "        :param x1: float，左上角 x 坐标\n",
    "        :param y1: float，左上角 y 坐标\n",
    "        :param x2: float，右下角 x 坐标\n",
    "        :param y2: float，右下角 y 坐标\n",
    "        \"\"\"\n",
    "        self.x1 = x1\n",
    "        self.y1 = y1\n",
    "        self.x2 = x2\n",
    "        self.y2 = y2\n",
    "\n",
    "    @classmethod\n",
    "    def from_xywh(cls, x, y, w, h):\n",
    "        \"\"\"\n",
    "        通过 xywh 格式创建边界框。\n",
    "        :param x: float，左上角 x\n",
    "        :param y: float，左上角 y\n",
    "        :param w: float，宽度\n",
    "        :param h: float，高度\n",
    "        :return: BBox 实例\n",
    "        \"\"\"\n",
    "        return cls(x, y, x + w, y + h)\n",
    "\n",
    "    @property\n",
    "    def xyxy(self):\n",
    "        \"\"\"返回 xyxy 格式：[x1, y1, x2, y2]\"\"\"\n",
    "        return [self.x1, self.y1, self.x2, self.y2]\n",
    "\n",
    "    @property\n",
    "    def xywh(self):\n",
    "        \"\"\"返回 xywh 格式：[x, y, w, h]\"\"\"\n",
    "        return [self.x1, self.y1, self.x2 - self.x1, self.y2 - self.y1]\n",
    "\n",
    "    def __repr__(self):\n",
    "        return f\"BBox(xyxy={self.xyxy}, xywh={self.xywh})\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "1bb28062-c378-4e55-9a17-800481e864ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "name2id = {\"xiantou\":2, \"zangwu\":1, \"yijiao\":0}\n",
    "id2name = {\"0\":\"yijiao\",\"1\":\"xiantou\",\"2\":\"zangwu\"}\n",
    "\n",
    "def parse_labels(ann_file, with_conf=False, task=\"segment\"):\n",
    "    results = []\n",
    "    with open(ann_file) as f:\n",
    "        annotations = f.readlines()\n",
    "        \n",
    "    for item in annotations:\n",
    "        item = item.split(\" \")\n",
    "        class_id = item[0]\n",
    "\n",
    "        if task == \"segment\":\n",
    "            if with_conf:\n",
    "                if len(item[1:]) < 2:\n",
    "                    continue\n",
    "                # print(item[1:-3])\n",
    "                # print(item)\n",
    "                coordinates = list(map(float, item[1:-1]))\n",
    "                score = float(item[-1])\n",
    "            else:\n",
    "                coordinates = list(map(float, item[1:]))\n",
    "                score = 1\n",
    "            polygon_points = [(coordinates[i], coordinates[i+1]) for i in range(0, len(coordinates), 2)]\n",
    "            x1, y1, x2, y2 = bbox_from_polygon(polygon_points)\n",
    "            bbox = BBox(x1, y1, x2, y2)\n",
    "        elif task == \"detect\":\n",
    "            if with_conf:\n",
    "                x,y,w,h = list(map(float, item[1:-1]))\n",
    "                x1, y1, x2, y2 = [x-w/2, y-h/2, x+w/2, y+h/2]\n",
    "                score = float(item[-1])\n",
    "                bbox = BBox(x1, y1, x2, y2)\n",
    "                polygon_points = None\n",
    "            else:\n",
    "                x,y,w,h = list(map(float, item[1:]))\n",
    "                x1, y1, x2, y2 = [x-w/2, y-h/2, x+w/2, y+h/2]\n",
    "                score = 1\n",
    "                bbox = BBox(x1, y1, x2, y2)\n",
    "                polygon_points = None\n",
    "                \n",
    "\n",
    "        flaw = Anomaly(score, bbox, polygon_points, class_id)\n",
    "        results.append(flaw)\n",
    "    return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "6e374433-d002-4b95-857c-86f31f153385",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1379/1379 [00:00<00:00, 1486.72it/s]\n"
     ]
    }
   ],
   "source": [
    "samples = []\n",
    "\n",
    "images = os.listdir(img_dir)\n",
    "for file in  tqdm(images):\n",
    "    ann_file = osp.join(anno_dir, file.replace(\".png\", \".txt\").replace(\".bmp\", \".txt\"))\n",
    "    pred_file = osp.join(preds_dir, file.replace(\".png\", \".txt\").replace(\".bmp\", \".txt\"))\n",
    "    annotations = parse_labels(ann_file, with_conf=False, task=\"detect\") \n",
    "\n",
    "    \n",
    "    real_path = osp.realpath(osp.join(img_dir, file))\n",
    "    sample = fo.Sample(filepath=real_path)\n",
    "    ground_truth = []\n",
    "    preds = []\n",
    "    for item in annotations:\n",
    "        ground_truth.append(fo.Detection(\n",
    "            label=item.category,\n",
    "            bounding_box=item.bbox.xywh\n",
    "        ))\n",
    "\n",
    "    if osp.exists(pred_file):\n",
    "        predictions = parse_labels(pred_file, with_conf=True, task=\"detect\")\n",
    "        for item in predictions:\n",
    "            preds.append(fo.Detection(\n",
    "                label=item.category,\n",
    "                bounding_box=item.bbox.xywh,\n",
    "                confidence=item.score\n",
    "            ))\n",
    "    sample[\"ground_truth\"] = fo.Detections(detections=ground_truth)\n",
    "    sample[\"predictions\"] = fo.Detections(detections=preds)\n",
    "    samples.append(sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "50e52d4e-cebc-4fa9-9db7-91543531f9f9",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 100% |███████████████| 1379/1379 [1.2s elapsed, 0s remaining, 1.1K samples/s]         \n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "        <iframe\n",
       "            width=\"100%\"\n",
       "            height=\"800\"\n",
       "            src=\"http://0.0.0.0:5151/?notebook=True&subscription=7ca7d649-5ef6-46a0-914b-cb6286634180\"\n",
       "            frameborder=\"0\"\n",
       "            allowfullscreen\n",
       "            \n",
       "        ></iframe>\n",
       "        "
      ],
      "text/plain": [
       "<IPython.lib.display.IFrame at 0x7fad94460dd0>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "dataset.add_samples(samples)\n",
    "\n",
    "session = fo.launch_app(dataset, address=\"0.0.0.0\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "fb3b2c8e-a732-47c2-9c6d-e068798491fd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Evaluating detections...\n",
      " 100% |███████████████| 1379/1379 [6.4s elapsed, 0s remaining, 217.3 samples/s]      \n",
      "Computing mistakenness...\n",
      " 100% |███████████████| 1379/1379 [7.0s elapsed, 0s remaining, 167.4 samples/s]       \n",
      "Mistakenness computation complete\n"
     ]
    }
   ],
   "source": [
    "fob.compute_mistakenness(dataset, \"predictions\", label_field=\"ground_truth\", mistakenness_field=\"mistakenness_loc\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "id": "9a214097-7405-4cd3-afe6-030eabe10501",
   "metadata": {},
   "outputs": [],
   "source": [
    "mistake_view = dataset.sort_by(\"mistakenness_loc\", reverse=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "id": "98cc8bb0-ae93-487c-94bd-d8cdc10b3d37",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset:     nike\n",
      "Media type:  image\n",
      "Num samples: 1897\n",
      "Sample fields:\n",
      "    id:                fiftyone.core.fields.ObjectIdField\n",
      "    filepath:          fiftyone.core.fields.StringField\n",
      "    tags:              fiftyone.core.fields.ListField(fiftyone.core.fields.StringField)\n",
      "    metadata:          fiftyone.core.fields.EmbeddedDocumentField(fiftyone.core.metadata.ImageMetadata)\n",
      "    ground_truth:      fiftyone.core.fields.EmbeddedDocumentField(fiftyone.core.labels.Detections)\n",
      "    predictions:       fiftyone.core.fields.EmbeddedDocumentField(fiftyone.core.labels.Detections)\n",
      "    mistakenness:      fiftyone.core.fields.IntField\n",
      "    possible_missing:  fiftyone.core.fields.IntField\n",
      "    possible_spurious: fiftyone.core.fields.IntField\n",
      "    mistakenness_loc:  fiftyone.core.fields.IntField\n",
      "View stages:\n",
      "    1. SortBy(field_or_expr='mistakenness_loc', reverse=True, create_index=True)\n"
     ]
    }
   ],
   "source": [
    "print(mistake_view)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "id": "e6436960-7e87-464f-93f7-729e77f28a92",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "        <iframe\n",
       "            width=\"100%\"\n",
       "            height=\"800\"\n",
       "            src=\"http://0.0.0.0:5151/?notebook=True&subscription=30a2c94c-ef1b-4f56-b870-90f01df2778d\"\n",
       "            frameborder=\"0\"\n",
       "            allowfullscreen\n",
       "            \n",
       "        ></iframe>\n",
       "        "
      ],
      "text/plain": [
       "<IPython.lib.display.IFrame at 0x7f5764a5cd70>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "session.view = mistake_view\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff85a922-83b5-4c32-996e-09db59404791",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
